package com.lap.mybatis.starter.plugins;

import com.lap.context.starter.handler.UserHandler;
import com.lap.framework.common.entity.Tenantable;
import java.lang.reflect.Array;
import java.util.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

@Slf4j
@Intercepts({
  @Signature(
      type = Executor.class,
      method = "update",
      args = {MappedStatement.class, Object.class}),
  @Signature(
      type = Executor.class,
      method = "query",
      args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class TenantInterceptor implements Interceptor {

  private static final String TENANT_CODE_KEY = "tenantCode";
  private static final Set<Class<?>> PRIMITIVE_WRAPPERS =
      Set.of(
          Byte.class,
          Short.class,
          Integer.class,
          Long.class,
          Float.class,
          Double.class,
          Character.class,
          Boolean.class,
          String.class);

  @Override
  public Object intercept(Invocation invocation) throws Throwable {
    Object[] args = invocation.getArgs();
    if (args == null || args.length < 2) return invocation.proceed();

    MappedStatement ms = (MappedStatement) args[0];
    Object parameter = args[1];
    String tenantCode = UserHandler.getTenantCode();
    if (StringUtils.isEmpty(tenantCode)) return invocation.proceed();

    args[1] = processParameter(ms, parameter, tenantCode);
    return invocation.proceed();
  }

  @SuppressWarnings("unchecked")
  private Object processParameter(MappedStatement ms, Object parameter, String tenantCode) {
    if (parameter == null || isPrimitiveOrWrapper(parameter.getClass())) {
      return wrapPrimitive(ms, parameter, tenantCode);
    } else if (parameter instanceof Map<?, ?> map) {
      Map<String, Object> m = (Map<String, Object>) map;
      m.putIfAbsent(TENANT_CODE_KEY, tenantCode);
      return m;
    } else if (parameter instanceof Collection<?> collection) {
      return processCollection(collection, tenantCode);
    } else if (parameter.getClass().isArray()) {
      return processArray(parameter, tenantCode);
    } else {
      setTenantIfPossible(parameter, tenantCode);
      return parameter;
    }
  }

  private Object wrapPrimitive(MappedStatement ms, Object value, String tenantCode) {
    MapperMethod.ParamMap<Object> map = new MapperMethod.ParamMap<>();
    map.put("param1", value);
    map.put(TENANT_CODE_KEY, tenantCode);

    try {
      BoundSql boundSql = ms.getBoundSql(value);
      List<ParameterMapping> pms = boundSql.getParameterMappings();
      if (pms != null) {
        for (ParameterMapping pm : pms) {
          String prop = pm.getProperty();
          if (!map.containsKey(prop)) {
            map.put(prop, value); // 用户自定义名字
          }
        }
      }
    } catch (Exception ignored) {
    }
    return map;
  }

  /** 处理集合，基础类型 / POJO 都支持 */
  private Object processCollection(Collection<?> collection, String tenantCode) {
    if (collection.isEmpty()) return collection;

    Object first = collection.iterator().next();
    if (first == null || isPrimitiveOrWrapper(first.getClass())) {
      return wrapIterable(collection, tenantCode);
    } else {
      collection.forEach(item -> setTenantIfPossible(item, tenantCode));
      return collection;
    }
  }

  /** 处理数组，基础类型 / POJO 都支持 */
  private Object processArray(Object array, String tenantCode) {
    int len = Array.getLength(array);
    if (len == 0) return array;

    Object first = Array.get(array, 0);
    if (first == null || isPrimitiveOrWrapper(first.getClass())) {
      List<Object> list = new ArrayList<>(len);
      for (int i = 0; i < len; i++) list.add(Array.get(array, i));
      return wrapIterable(list, tenantCode);
    } else {
      for (int i = 0; i < len; i++) setTenantIfPossible(Array.get(array, i), tenantCode);
      return array;
    }
  }

  /** 包装基础类型集合 / 数组的 ParamMap */
  private MapperMethod.ParamMap<Object> wrapIterable(Object iterable, String tenantCode) {
    MapperMethod.ParamMap<Object> map = new MapperMethod.ParamMap<>();
    map.put("param1", iterable);
    map.put("list", iterable);
    map.put("collection", iterable);
    map.put(TENANT_CODE_KEY, tenantCode);
    return map;
  }

  private void setTenantIfPossible(Object obj, String tenantCode) {
    if (obj instanceof Tenantable tenantable && tenantCode != null) {
      tenantable.setTenantCode(tenantCode);
    }
  }

  private boolean isPrimitiveOrWrapper(Class<?> clazz) {
    return clazz.isPrimitive() || PRIMITIVE_WRAPPERS.contains(clazz);
  }

  @Override
  public Object plugin(Object target) {
    return Plugin.wrap(target, this);
  }

  @Override
  public void setProperties(Properties properties) {
    log.info("Initializing TenantInterceptor with properties: {}", properties);
  }
}
